import pandas as pd
# BertForSequenceClassification 包括了最终的二分类模型数据
from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
# 之后，你可以使用Transformers库的`AutoModel.from_pretrained`来加载这个模型
from transformers import AutoModel


# 加载数据
train_file = "D:/code/datasets/train_v2_drcat_02/train_v2_drcat_02.csv"
# 加载数据
df = pd.read_csv(train_file, usecols=['text', 'label'])

# 文本预处理函数
# def preprocess_function(examples):
#     return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)

# 初始化tokenizer和模型
print("开始初始化bert模型")
# tokenizer = BertTokenizer.from_pretrained('D:/code/models/bert_classsfication')
model = BertForSequenceClassification.from_pretrained('D:/code/models/bert_classsfication')
# from transformers import AutoModel
# loaded_model = AutoModel.from_pretrained('D:/code/models/bert_classsfication')

trainer = Trainer(
    model=model
)
print(model)
# # 预测验证集
# predictions = trainer.predict(val_dataset)